#include "gtest/gtest.h"
#include <mockcpp/mokc.h>
#include <mockcpp/mockcpp.hpp>

#include <vector>
#include <iostream>

#include "topoinfo_struct.h"
#include "log.h"
#include "checker.h"
#include "virtual_topo.h"
#include "testcase_utils.h"

using namespace hccl;

class ReduceTest : public testing::Test {
protected:
    static void SetUpTestCase()
    {
        std::cout << "ReduceTest set up." << std::endl;
    }

    static void TearDownTestCase()
    {
        std::cout << "ReduceTest tear down." << std::endl;
    }

    virtual void SetUp()
    {
        const ::testing::TestInfo* const test_info = ::testing::UnitTest::GetInstance()->current_test_info();
        std::string caseName = "analysis_result_" + std::string(test_info->test_case_name()) + "_" + std::string(test_info->name());
        Checker::SetDumpFileName(caseName);
    }

    virtual void TearDown()
    {
        Checker::SetDumpFileName("analysis_result");
        GlobalMockObject::verify();
        // 这边每个case执行完成需要清理所有的环境变量，如果有新增的环境变量，需要在这个函数中进行清理
        ClearHcclEnv();
    }
};

TEST_F(ReduceTest, reduce_test_opbase_910B_loop)
{
    RankTable_For_LLT gen;
    TestOpParam testOpParam;
    testOpParam.opType = HcclCMDType::HCCL_CMD_REDUCE;
    testOpParam.tag = "AllReduce";
    testOpParam.opMode = OPBASE;
    testOpParam.reduceType = HcclReduceOp::HCCL_REDUCE_SUM;
    testOpParam.devtype = DevType::DEV_TYPE_910B;
    testOpParam.DataDes.count = 100;
    testOpParam.DataDes.dataType = HcclDataType::HCCL_DATA_TYPE_INT32;
    testOpParam.root = 0;

    HcclResult ret;
    for (int i = 1; i <= 2; i++) {
        for (int j = 1; j <= 8; j++) {
            Checker checker;
            RankTable_t ranktable;
            TopoMeta topoMeta;
            gen.GenTopoMeta(topoMeta, 1, i, j);
            gen.GenRankTable(topoMeta, ranktable);
            checker.CloseRankMemCheck();
            ret = checker.Check(testOpParam, ranktable, topoMeta);
            EXPECT_EQ(ret, HcclResult::HCCL_SUCCESS);
        }
    }
}

TEST_F(ReduceTest, reduce_test_offload_910B_loop)
{
    RankTable_For_LLT gen;
    TestOpParam testOpParam;
    testOpParam.opType = HcclCMDType::HCCL_CMD_REDUCE;
    testOpParam.tag = "Reduce";
    testOpParam.opMode = OFFLOAD;
    testOpParam.reduceType = HcclReduceOp::HCCL_REDUCE_SUM;
    testOpParam.devtype = DevType::DEV_TYPE_910B;
    testOpParam.DataDes.count = 100;
    testOpParam.DataDes.dataType = HcclDataType::HCCL_DATA_TYPE_INT32;
    testOpParam.root = 0;

    HcclResult ret;
    for (int i = 1; i <= 2; i++) {
        for (int j = 1; j <= 8; j++) {
            Checker checker;
            RankTable_t ranktable;
            TopoMeta topoMeta;
            gen.GenTopoMeta(topoMeta, 1, i, j);
            gen.GenRankTable(topoMeta, ranktable);
            checker.CloseRankMemCheck();
            ret = checker.Check(testOpParam, ranktable, topoMeta);
            EXPECT_EQ(ret, HcclResult::HCCL_SUCCESS);
        }
    }
}

TEST_F(ReduceTest, reduce_test_offload_910A)
{
    RankTable_For_LLT gen;
    TestOpParam testOpParam;
    testOpParam.opType = HcclCMDType::HCCL_CMD_REDUCE;
    testOpParam.tag = "Reduce";
    testOpParam.opMode = OFFLOAD;
    testOpParam.reduceType = HcclReduceOp::HCCL_REDUCE_SUM;
    testOpParam.devtype = DevType::DEV_TYPE_910;
    testOpParam.DataDes.count = 100;
    testOpParam.DataDes.dataType = HcclDataType::HCCL_DATA_TYPE_INT32;
    testOpParam.root = 0;

    HcclResult ret;

    Checker checker;
    RankTable_t ranktable;
    TopoMeta topoMeta;
    gen.GenTopoMeta(topoMeta, 1, 1, 8);
    gen.GenRankTable(topoMeta, ranktable);
    ret = checker.Check(testOpParam, ranktable, topoMeta);
    EXPECT_EQ(ret, HcclResult::HCCL_SUCCESS);
}

TEST_F(ReduceTest, reduce_test_opbase_910A)
{
    RankTable_For_LLT gen;
    TestOpParam testOpParam;
    testOpParam.opType = HcclCMDType::HCCL_CMD_REDUCE;
    testOpParam.tag = "Reduce";
    testOpParam.opMode = OPBASE;
    testOpParam.reduceType = HcclReduceOp::HCCL_REDUCE_SUM;
    testOpParam.devtype = DevType::DEV_TYPE_910;
    testOpParam.DataDes.count = 100;
    testOpParam.DataDes.dataType = HcclDataType::HCCL_DATA_TYPE_INT32;
    testOpParam.root = 0;

    HcclResult ret;
    Checker checker;
    RankTable_t ranktable;
    TopoMeta topoMeta;
    gen.GenTopoMeta(topoMeta, 1, 1, 8);
    gen.GenRankTable(topoMeta, ranktable);
    ret = checker.Check(testOpParam, ranktable, topoMeta);
    EXPECT_EQ(ret, HcclResult::HCCL_SUCCESS);
}

TEST_F(ReduceTest, reduce_test_opbase_910_93)
{
    RankTable_For_LLT gen;
    TestOpParam testOpParam;
    testOpParam.opType = HcclCMDType::HCCL_CMD_REDUCE;
    testOpParam.tag = "Reduce";
    testOpParam.opMode = OPBASE;
    testOpParam.reduceType = HcclReduceOp::HCCL_REDUCE_SUM;
    testOpParam.devtype = DevType::DEV_TYPE_910_93;
    testOpParam.DataDes.count = 100;
    testOpParam.DataDes.dataType = HcclDataType::HCCL_DATA_TYPE_INT32;
    testOpParam.root = 0;

    HcclResult ret;
    Checker checker;
    RankTable_t ranktable;
    TopoMeta topoMeta;
    gen.GenTopoMeta(topoMeta, 1, 1, 8);
    gen.GenRankTable(topoMeta, ranktable);
    ret = checker.Check(testOpParam, ranktable, topoMeta);
    EXPECT_EQ(ret, HcclResult::HCCL_SUCCESS);
}

TEST_F(ReduceTest, reduce_test_opbase_offload_93_2pods)
{
    RankTable_For_LLT gen;
    TestOpParam testOpParam;
    testOpParam.opType = HcclCMDType::HCCL_CMD_REDUCE;
    testOpParam.tag = "Reduce";
    testOpParam.opMode = OFFLOAD;
    testOpParam.reduceType = HcclReduceOp::HCCL_REDUCE_SUM;
    testOpParam.devtype = DevType::DEV_TYPE_910_93;
    testOpParam.DataDes.count = 100;
    testOpParam.DataDes.dataType = HcclDataType::HCCL_DATA_TYPE_INT32;
    testOpParam.root = 0;

    HcclResult ret;
    Checker checker;
    RankTable_t ranktable;
    TopoMeta topoMeta;
    gen.GenTopoMeta(topoMeta, 2, 2, 8);
    gen.GenRankTable(topoMeta, ranktable);
    ret = checker.Check(testOpParam, ranktable, topoMeta);
    EXPECT_EQ(ret, HcclResult::HCCL_SUCCESS);
}

TEST_F(ReduceTest, reduce_test_opbase_ReduceComm)
{
    RankTable_For_LLT gen;
    TestOpParam testOpParam;
    testOpParam.opType = HcclCMDType::HCCL_CMD_REDUCE;
    testOpParam.tag = "Reduce";
    testOpParam.algName = "ReduceComm";
    testOpParam.opMode = OFFLOAD;
    testOpParam.reduceType = HcclReduceOp::HCCL_REDUCE_SUM;
    testOpParam.devtype = DevType::DEV_TYPE_910B;
    testOpParam.DataDes.count = 100;
    testOpParam.DataDes.dataType = HcclDataType::HCCL_DATA_TYPE_INT32;
    testOpParam.root = 0;

    HcclResult ret;
    Checker checker;
    RankTable_t ranktable;
    TopoMeta topoMeta;
    gen.GenTopoMeta(topoMeta, 1, 1, 8);
    gen.GenRankTable(topoMeta, ranktable);
    ret = checker.Check(testOpParam, ranktable, topoMeta);
    EXPECT_EQ(ret, HcclResult::HCCL_SUCCESS);
}

TEST_F(ReduceTest, reduce_test_opbase_ReduceMeshExecutor)
{
    RankTable_For_LLT gen;
    TestOpParam testOpParam;
    testOpParam.opType = HcclCMDType::HCCL_CMD_REDUCE;
    testOpParam.tag = "Reduce";
    testOpParam.algName = "ReduceMeshExecutor";
    testOpParam.opMode = OFFLOAD;
    testOpParam.reduceType = HcclReduceOp::HCCL_REDUCE_SUM;
    testOpParam.devtype = DevType::DEV_TYPE_910B;
    testOpParam.DataDes.count = 100;
    testOpParam.DataDes.dataType = HcclDataType::HCCL_DATA_TYPE_INT32;
    testOpParam.root = 0;

    HcclResult ret;
    Checker checker;
    RankTable_t ranktable;
    TopoMeta topoMeta;
    gen.GenTopoMeta(topoMeta, 1, 1, 8);
    gen.GenRankTable(topoMeta, ranktable);
    checker.CloseRankMemCheck();
    ret = checker.Check(testOpParam, ranktable, topoMeta);
    EXPECT_EQ(ret, HcclResult::HCCL_SUCCESS);
}

TEST_F(ReduceTest, reduce_test_opbase_ReduceRingPlusHd)
{
    RankTable_For_LLT gen;
    TestOpParam testOpParam;
    testOpParam.opType = HcclCMDType::HCCL_CMD_REDUCE;
    testOpParam.tag = "Reduce";
    testOpParam.algName = "ReduceRingPlusHd";
    testOpParam.opMode = OFFLOAD;
    testOpParam.reduceType = HcclReduceOp::HCCL_REDUCE_SUM;
    testOpParam.devtype = DevType::DEV_TYPE_910B;
    testOpParam.DataDes.count = 100;
    testOpParam.DataDes.dataType = HcclDataType::HCCL_DATA_TYPE_INT32;
    testOpParam.root = 0;

    HcclResult ret;
    Checker checker;
    RankTable_t ranktable;
    TopoMeta topoMeta;
    gen.GenTopoMeta(topoMeta, 1, 1, 8);
    gen.GenRankTable(topoMeta, ranktable);
    ret = checker.Check(testOpParam, ranktable, topoMeta);
    EXPECT_EQ(ret, HcclResult::HCCL_SUCCESS);
}

TEST_F(ReduceTest, reduce_test_opbase_ReduceSingleExecutor)
{
    RankTable_For_LLT gen;
    TestOpParam testOpParam;
    testOpParam.opType = HcclCMDType::HCCL_CMD_REDUCE;
    testOpParam.tag = "Reduce";
    testOpParam.algName = "ReduceSingleExecutor";
    testOpParam.opMode = OFFLOAD;
    testOpParam.reduceType = HcclReduceOp::HCCL_REDUCE_SUM;
    testOpParam.devtype = DevType::DEV_TYPE_910B;
    testOpParam.DataDes.count = 100;
    testOpParam.DataDes.dataType = HcclDataType::HCCL_DATA_TYPE_INT32;
    testOpParam.root = 0;

    HcclResult ret;
    Checker checker;
    RankTable_t ranktable;
    TopoMeta topoMeta;
    gen.GenTopoMeta(topoMeta, 1, 1, 1);
    gen.GenRankTable(topoMeta, ranktable);
    ret = checker.Check(testOpParam, ranktable, topoMeta);
    EXPECT_EQ(ret, HcclResult::HCCL_SUCCESS);
}

TEST_F(ReduceTest, reduce_test_opbase_910_93_2superpod_ReduceMeshExecutor_postSync)
{
    RankTable_For_LLT gen;
    TestOpParam testOpParam;
    testOpParam.opType = HcclCMDType::HCCL_CMD_REDUCE;
    testOpParam.tag = "Reduce";
    testOpParam.opMode = OPBASE;
    testOpParam.reduceType = HcclReduceOp::HCCL_REDUCE_SUM;
    testOpParam.devtype = DevType::DEV_TYPE_910_93;
    testOpParam.DataDes.count = 100;
    testOpParam.DataDes.dataType = HcclDataType::HCCL_DATA_TYPE_INT32;
    testOpParam.root = 0;
    testOpParam.aicpuUnfoldMode = true;
    testOpParam.algOpContext.opRetryHandler.isPostSync = true;

    HcclResult ret;
    Checker checker;
    RankTable_t ranktable;
    TopoMeta topoMeta;
    gen.GenTopoMeta(topoMeta, 2, 1, 1);
    gen.GenRankTable(topoMeta, ranktable);
    ret = checker.Check(testOpParam, ranktable, topoMeta);

}

TEST_F(ReduceTest, reduce_test_opbase_910_93_2superpod_4servers_4ranks_Reduce)
{
    RankTable_For_LLT gen;
    RankTable_t ranktable;
    TopoMeta topoMeta;
    gen.GenTopoMeta(topoMeta, 2, 4, 4);
    gen.GenRankTable(topoMeta, ranktable);

    setenv("HCCL_INTER_HCCS_DISABLE", "TRUE", 1);
    setenv("HCCL_ALGO", "level1:ring", 1);

    TestOpParam testOpParam;
    testOpParam.opType = HcclCMDType::HCCL_CMD_REDUCE;
    testOpParam.tag = "Reduce";
    testOpParam.opMode = OPBASE;
    testOpParam.reduceType = HcclReduceOp::HCCL_REDUCE_SUM;
    testOpParam.devtype = DevType::DEV_TYPE_910_93;
    testOpParam.DataDes.count = 8192;
    testOpParam.DataDes.dataType = HcclDataType::HCCL_DATA_TYPE_FP32;
    testOpParam.root = 0;

    HcclResult ret;
    Checker checker;
    
    ret = checker.Check(testOpParam, ranktable, topoMeta);
    EXPECT_EQ(ret, HcclResult::HCCL_SUCCESS);
}